package com.shujia.flink.state;

import lombok.AllArgsConstructor;
import lombok.Data;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.api.common.serialization.SimpleStringSchema;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.connector.kafka.source.KafkaSource;
import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.util.Collector;

/**
 * @author shujia
 */
public class Demo6Timer {
    public static void main(String[] args) throws Exception {

        //创建环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        //读取数据
        KafkaSource<String> kafkaSource = KafkaSource.<String>builder()
                //kafka broker列表
                .setBootstrapServers("master:9092,node1:9092,node2:9092")
                //指定topic
                .setTopics("lines")
                //消费者组
                .setGroupId("Demo1NoState")
                //读取数据的位置
                .setStartingOffsets(OffsetsInitializer.latest())
                //数据格式
                .setValueOnlyDeserializer(new SimpleStringSchema())
                .build();

        //使用kafka source
        DataStream<String> lines = env
                .fromSource(kafkaSource, WatermarkStrategy.noWatermarks(), "kafka source");

        //一行转换成多行
        DataStream<String> words = lines.flatMap((line, collect) -> {
            for (String word : line.split(",")) {
                collect.collect(word);
            }
        }, Types.STRING);

        //分组
        KeyedStream<String, String> keyBys = words.keyBy(word -> word);

        /*
         * 通过状态+定时器时间会话窗口的功能
         * 会话窗口：一段时间没有数据开始计算
         */

        DataStream<WordCount> wordCounts = keyBys.process(new KeyedProcessFunction<String, String, WordCount>() {

            ValueState<Long> timerState;

            ListState<String> wordsState;

            @Override
            public void open(Configuration parameters) throws Exception {
                RuntimeContext context = getRuntimeContext();

                //用于保存定时器时间的状态
                timerState = context.getState(new ValueStateDescriptor<>("timer", Types.LONG));

                //用于保存所有单词的状态
                wordsState = context.getListState(new ListStateDescriptor<>("words", Types.STRING));
            }

            @Override
            public void processElement(String word,
                                       KeyedProcessFunction<String, String, WordCount>.Context ctx,
                                       Collector<WordCount> out) throws Exception {

                //从状态中获取定时器
                Long ts = timerState.value();
                if (ts != null) {
                    //删除定时器
                    ctx.timerService().deleteProcessingTimeTimer(ts);
                }

                //获取当前时间戳
                Long timestamp = ctx.timestamp();
                System.out.println("正在处理时间，注册定时器：" + timestamp);
                //注册定时器
                //timestamp + 5000： 在5秒之后执行onTimer方法
                ctx.timerService().registerProcessingTimeTimer(timestamp + 5000);
                //将定时器二点时间保存到状态中
                timerState.update(timestamp + 5000);


                //将单词保存到状态中
                wordsState.add(word);
            }

            /**
             * onTimer方法是定时器触发之后执行的方法
             * @param timestamp: 触发的时间
             * @param ctx 上下文对象
             * @param out 用于将数据发送到下游
             */
            @Override
            public void onTimer(long timestamp,
                                KeyedProcessFunction<String, String, WordCount>.OnTimerContext ctx,
                                Collector<WordCount> out) throws Exception {

                //当前当前的key
                String word = ctx.getCurrentKey();

                System.out.println("定时器被触发：" + timestamp);

                //统计单词的数量
                Iterable<String> words = wordsState.get();
                int count = 0;
                for (String w : words) {
                    count++;
                }

                //将计算结果发送到下游
                out.collect(new WordCount(word, count));

                //清除状态
                wordsState.clear();
            }
        });

        wordCounts.print();

        env.execute();


    }

    @Data
    @AllArgsConstructor
    static class WordCount {
        private String word;
        private Integer count;
    }
}
